package cn.itcast.flink.base;

import org.apache.flink.api.common.restartstrategy.RestartStrategies;
import org.apache.flink.api.common.state.BroadcastState;
import org.apache.flink.api.common.state.MapStateDescriptor;
import org.apache.flink.api.common.state.ReadOnlyBroadcastState;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple4;
import org.apache.flink.api.java.tuple.Tuple6;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.BroadcastConnectedStream;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.util.Collector;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.text.SimpleDateFormat;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;

/**
 * Author itcast
 * Date 2021/7/31 7:47
 * Desc TODO
 */
public class BroadcastStateUpdateDemo {
    public static void main(String[] args) throws Exception {
        //1.创建流执行环境，设置并行度，设置状态后端设置和属性，重启策略
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        env.setRestartStrategy(RestartStrategies.fixedDelayRestart(3, 3000));
        //2.自定义数据源
        //2-1.自定义实时数据流，设置随机数据并生成 Tuple4<String, String, String, Integer> 类型
        DataStreamSource<Tuple4<String, String, String, Integer>> source = env.addSource(new MySource());
        //2-2.从MySQL中读取用户的配置数据 Map<String/*用户id*/,Tuple2<String/*姓名*/,Integer/*年龄*/>>
        DataStreamSource<HashMap<String, Tuple2<String, Integer>>> userConfig = env.addSource(new UserConfigInfo());
        //3.transformation
        //3-1.定义状态描述器-将配置流作为状态广播
        //设置描述器descriptor = new MapStateDescriptor<Void, Map<String, Tuple2<String, Integer>>>("config",Types.VOID, Types.MAP(Types.STRING, Types.TUPLE(Types.STRING, Types.INT)));
        MapStateDescriptor<Void, Map<String, Tuple2<String, Integer>>> desc = new MapStateDescriptor<>(
                "config",
                Types.VOID,
                Types.MAP(Types.STRING, Types.TUPLE(Types.STRING, Types.INT))
        );
        //3-2.将用户配置数据源广播出去 configDS.broadcast(descriptor);
        BroadcastStream<HashMap<String, Tuple2<String, Integer>>> broadcast = userConfig.broadcast(desc);
        //3-3.将用户数据流connect用户配置广播流
        BroadcastConnectedStream<Tuple4<String, String, String, Integer>, HashMap<String, Tuple2<String, Integer>>> connectDataStream = source
                .connect(broadcast);
        //3-4.处理(process)连接后的流-根据配置流补全事件流中的用户的信息,Tuple4和Tuple2合并
        SingleOutputStreamOperator<Tuple6<String, String, String, Integer, String, Integer>> processDataStream = connectDataStream.process(new BroadcastProcessFunction<Tuple4<String, String, String, Integer>,
                HashMap<String, Tuple2<String, Integer>>, Tuple6<String, String, String, Integer, String, Integer>>() {
            //3-4-1.processElement 将广播变量中的用户配置信息和用户数据流合并成Tuple6<String, String, String, Integer, String, Integer>
            @Override
            public void processElement(Tuple4<String, String, String, Integer> value, ReadOnlyContext ctx,
                                       Collector<Tuple6<String, String, String, Integer, String, Integer>> out) throws Exception {
                if (value != null) {
                    //获取当前用户信息
                    String userId = value.f0;
                    //获取广播变量信息
                    ReadOnlyBroadcastState<Void, Map<String, Tuple2<String, Integer>>> broadcastState = ctx.getBroadcastState(desc);
                    if (broadcastState != null) {
                        Map<String, Tuple2<String, Integer>> userConfigInfo = broadcastState.get(null);
                        if (userConfigInfo != null) {
                            Tuple2<String, Integer> userNameAge = userConfigInfo.get(userId);
                            if (userNameAge != null) {
                                String username = userNameAge.f0;
                                int age = userNameAge.f1;
                                out.collect(Tuple6.of(userId, value.f1, value.f2, value.f3, username, age));
                            }
                        }
                    }
                }
            }

            //3-4-2.processBroadcastElement 将广播变量状态里的用户配置数据进行更细
            @Override
            public void processBroadcastElement(HashMap<String, Tuple2<String, Integer>> value, Context ctx,
                                                Collector<Tuple6<String, String, String, Integer, String, Integer>> out) throws Exception {
                //3.4.2.1.获取内存中的广播变量数据
                BroadcastState<Void, Map<String, Tuple2<String, Integer>>> broadcastState = ctx.getBroadcastState(desc);
                //3.4.2.2.在清空历史状态数据
                broadcastState.clear();
                //3.4.2.3最后将最新的广播流数据放到state中（更新状态数据）
                broadcastState.put(null, value);
            }
        });
        //4.sinks
        processDataStream.print();
        //5.execute
        env.execute();
    }

    private static class MySource implements SourceFunction<Tuple4<String, String, String, Integer>> {
        boolean flag = true;
        Random rn = new Random();
        SimpleDateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
        String[] eventTypes = {"browser", "click"};

        @Override
        public void run(SourceContext<Tuple4<String, String, String, Integer>> ctx) throws Exception {
            while (flag) {
                int id = rn.nextInt(4);
                String userID = "user_" + id;
                String eventTime = sdf.format(new Date());
                String eventType = eventTypes[rn.nextInt(2)];
                int productID = rn.nextInt(4);
                ctx.collect(Tuple4.of(userID, eventTime, eventType, productID));
                Thread.sleep(1000);
            }
        }

        @Override
        public void cancel() {
            flag = false;
        }
    }

    private static class UserConfigInfo extends RichSourceFunction<HashMap<String, Tuple2<String, Integer>>> {
        Connection conn = null;
        PreparedStatement ps = null;
        boolean flag = true;
        ResultSet rs = null;

        @Override
        public void open(Configuration parameters) throws Exception {
            conn = DriverManager.getConnection("jdbc:mysql://node3:3306/bigdata?useSSL=false", "root", "123456");
            String sql = "select userID,userName,userAge from user_info";
            ps = conn.prepareStatement(sql);
        }

        @Override
        public void run(SourceContext<HashMap<String, Tuple2<String, Integer>>> ctx) throws Exception {
            HashMap<String, Tuple2<String, Integer>> userConfig = new HashMap<>();
            rs = ps.executeQuery();
            while (flag) {
                while (rs.next()) {
                    String userID = rs.getString("userID");
                    String userName = rs.getString("userName");
                    int userAge = rs.getInt("userAge");
                    userConfig.put(userID, Tuple2.of(userName, userAge));
                }
                ctx.collect(userConfig);
                Thread.sleep(5000L);
            }
        }

        @Override
        public void cancel() {
            flag = false;
        }

        @Override
        public void close() throws Exception {
            if (!rs.isClosed()) rs.close();
            if (!ps.isClosed()) ps.close();
            if (!conn.isClosed()) conn.close();
        }
    }
}
